#include <planner/distance_map_planner/distance_map_planner.h>
#include "map/grid_map.h"
namespace sz {

DMPlanner::DMPlanner(bool verbose) : planner_verbose_(verbose) {
  planner_verbose_ = verbose;
  if (planner_verbose_)
    printf(ANSI_COLOR_CYAN "DMP PLANNER VERBOSE ON\n" ANSI_COLOR_RESET);
}

void DMPlanner::SetSearchRadius(const Point2f &r) {
  search_radius_ = r;
}

void DMPlanner::SetPotentialRadius(const Point2f &r) {
  potential_radius_ = r;
}

__attribute__((unused)) void DMPlanner::SetPotentialMapRange(const Point2f &r) {
  potential_map_range_ = r;
}

__attribute__((unused)) void DMPlanner::SetEps(float eps) {
  eps_ = eps;
}

__attribute__((unused)) void DMPlanner::SetCollisionWeight(float c) {
  cweight_ = c;
}

void DMPlanner::SetPow(int pow) {
  if (pow < 1) {
    std::cerr << "Too little pow" << std::endl;
    return;
  }
  pow_ = pow;
}

int DMPlanner::Status() const { return status_; }

VectorPoint2f DMPlanner::GetRawPath() { return raw_path_; }
VectorPoint2i DMPlanner::GetMapRawPath() { return map_raw_path_; }

__attribute__((unused)) VectorPoint2i DMPlanner::GetPriorPath() { return prior_path_; }

__attribute__((unused)) VectorPoint2f DMPlanner::GetOpenSet() const {
  VectorPoint2f ps;
  const auto ss = graph_search_->GetOpenSet();
  for (const auto &it: ss) {
    ps.push_back(map_util_->MapToWorld(Point2i(it->x, it->y)));
  }
  return ps;
}

__attribute__((unused)) VectorPoint2f DMPlanner::GetCloseSet() const {
  VectorPoint2f ps;
  const auto ss = graph_search_->GetCloseSet();
  for (const auto &it: ss) {
    ps.push_back(map_util_->MapToWorld(Point2i(it->x, it->y)));
  }
  return ps;
}

__attribute__((unused)) VectorPoint2f DMPlanner::GetAllSet() const {
  VectorPoint2f ps;
  const auto ss = graph_search_->GetAllSet();
  for (const auto &it: ss) {
    ps.push_back(map_util_->MapToWorld(Point2i(it->x, it->y)));
  }
  return ps;
}

std::vector<bool> DMPlanner::SetPath(const VectorPoint2i &path,
                                     const Point2f &radius, bool dense) {
  prior_path_ = path;
  // create cells along path
  vec_Veci<2> ps;
  if (!dense) {
    for (unsigned int i = 1; i < path.size(); i++) {
      auto pns = map_util_->RayTrace(path[i - 1], path[i]);
      //ps.push_back(map_util_->floatToInt(path[i-1]));
      ps.insert(ps.end(), pns.begin(), pns.end());
      ps.push_back(path[i]);
    }
  } else {
    for (const auto &pt: path)
      ps.push_back(pt);
  }

  // create in_region zslam_map
  auto dim = map_util_->GetDim();
  std::vector<bool> in_region;
  // if radius is negative, set no tunnel
  if (radius(0) < 0 || radius(1) < 0) {

    in_region.resize(dim(0) * dim(1), true);

    return in_region;
  }
  in_region.resize(dim(0) * dim(1), false);

  // create mask
  vec_Veci<2> ns;
  int rn = std::ceil(radius(0) / map_util_->GetRes());

  for (int nx = -rn; nx <= rn; nx++) {
    for (int ny = -rn; ny <= rn; ny++) {
      if (std::hypot(nx, ny) > rn)
        continue;
      ns.emplace_back(nx, ny);
    }
  }

  for (const auto &it: ps) {
    for (const auto &n: ns) {
      Point2i pn = it + n;
      if (map_util_->IsOutside(pn))
        continue;
      int idx = pn(0) + dim(0) * pn(1);
      if (!in_region[idx]) {
        in_region[idx] = true;
      }
    }

  }

  return in_region;
}

__attribute__((unused)) VectorPoint2f DMPlanner::GetSearchRegion() {
  auto dim = map_util_->GetDim();
  VectorPoint2f pts;

  for (int nx = 0; nx < dim(0); nx++) {
    for (int ny = 0; ny < dim(1); ny++) {
      int idx = nx + dim(0) * ny;
      if (search_region_[idx])
        pts.push_back(map_util_->MapToWorld(Point2i(nx, ny)));
    }
  }

  return pts;
}

__attribute__((unused)) std::shared_ptr<sz::NavigationMap> DMPlanner::GetMapUtil() {
  return map_util_;
}

bool DMPlanner::CheckAvailability(const Point2i &pn) {
  if (map_util_->IsSLAMOccupied(pn)) {
    if (planner_verbose_)
      printf(ANSI_COLOR_RED "point is unknown!\n" ANSI_COLOR_RESET);
    return false;
  }
  if (map_util_->IsOutside(pn)) {
    if (planner_verbose_) {
      printf(ANSI_COLOR_RED "point is outside!\n" ANSI_COLOR_RESET);
      std::cout << "coordinate: " << pn.transpose() << std::endl;
      std::cout << "Map origin: " << map_util_->GetMax().transpose() << std::endl;
      std::cout << "Map dim: " << map_util_->GetDim().transpose() << std::endl;
    }
    return false;
  }
  if (cmap_[map_util_->GetIndex(pn)] == H_MAX) {
    if (planner_verbose_)
      printf(ANSI_COLOR_RED "point is occupied!\n" ANSI_COLOR_RESET);
    return false;
  }
  return true;
}

VectorEigen<std::pair<Point2i, uint8_t>> DMPlanner::CreateMask(int pow) {
  /// Mask for generating potential field around obstacle
  VectorEigen<std::pair<Point2i, uint8_t>> mask;
  // create mask
  float res = map_util_->GetRes();
  float h_max = H_MAX;
  short rn = std::ceil(potential_radius_(0) / res);
  Point2i n;

  for (n(0) = (short) -rn; n(0) <= rn; n(0)++) {
    for (n(1) = (short) -rn; n(1) <= rn; n(1)++) {
      if (std::hypot(n(0), n(1)) > rn)
        continue;
      auto h = h_max * std::pow<float>((float) (1.f - std::hypot((float) n(0), (float) n(1)) / (float) rn), pow);
      if (h > 1e-3)
        mask.emplace_back(std::make_pair(n, (int8_t) h));
    }
  }

  return mask;
}

void DMPlanner::SetMap(const std::shared_ptr<sz::NavigationMap> &map_util,
                       const Point2f &pos) {
  map_util_ = map_util;
  const auto mask = CreateMask(pow_);
  // compute a 2D local distance zslam_map
  const auto dim = map_util_->GetDim();
  Point2i coord1 = Point2i::Zero();
  Point2i coord2 = dim;
  if (potential_map_range_.norm() > 0) {
    coord1 = map_util_->WorldToMap(pos - potential_map_range_);
    coord2 = map_util_->WorldToMap(pos + potential_map_range_);
    for (int i = 0; i < 2; i++) {
      if (coord1(i) < 0)
        coord1(i) = 0;
      else if (coord1(i) >= dim(i))
        coord1(i) = short(dim(i) - 1);

      if (coord2(i) < 0)
        coord2(i) = 0;
      else if (coord2(i) >= dim(i))
        coord2(i) = short(dim(i) - 1);
    }
  }

  const std::vector<uint8_t> &map = map_util_->GetMap();
  auto distance_map = map;

  Point2i n;

  for (n(0) = coord1(0); n(0) < coord2(0); n(0)++) {
    for (n(1) = coord1(1); n(1) < coord2(1); n(1)++) {
      int idx = map_util_->GetIndex(n);
      if (map[idx] == 0) {
        continue;
      }
      distance_map[idx] = H_MAX;
      for (const auto &it: mask) {
        const Point2i new_n = n + it.first;

        if (!map_util_->IsOutside(new_n)) {
          const int new_idx = map_util_->GetIndex(new_n);
          distance_map[new_idx] = std::max(distance_map[new_idx], it.second);
        }
      }
    }
  }

  cmap_ = distance_map;
  map_util_->SetMap(map_util_->GetMapInfo(), distance_map);
}

bool DMPlanner::Plan(const Point2f &start, const Point2f &goal, float_decimal_t eps, float_decimal_t cweight) {
  if (planner_verbose_) {
    std::cout << "Start: " << start.transpose() << std::endl;
    std::cout << "Goal:  " << goal.transpose() << std::endl;
    std::cout << "Epsilon:  " << eps << std::endl;
    std::cout << "cweight:  " << cweight << std::endl;
  }
  VectorPoint2f().swap(raw_path_);
  status_ = 0;

  /// check if the zslam_map exists
  if (cmap_.empty()) {
    if (planner_verbose_)
      printf(ANSI_COLOR_RED "need to set zslam_map, call SetMap()!\n" ANSI_COLOR_RESET);
    status_ = -1;
    return false;
  }


  /// check availability of start
  const Point2i start_int = map_util_->WorldToMap(start);
  if (!CheckAvailability(start_int)) {
    if (planner_verbose_)
      printf(ANSI_COLOR_RED "start is invalid!\n" ANSI_COLOR_RESET);
    status_ = 1;
    return false;
  }

  /// check availability of goal
  const Point2i goal_int = map_util_->WorldToMap(goal);
  if (!CheckAvailability(goal_int)) {
    if (planner_verbose_)
      printf(ANSI_COLOR_RED "goal is invalid!\n" ANSI_COLOR_RESET);
    status_ = 2;
    return false;
  }

  const Point2i dim = map_util_->GetDim();

  graph_search_ = std::make_shared<DMPGraphSearch>(
      cmap_.data(), dim(0), dim(1), eps, cweight, planner_verbose_);
  path_cost_ = graph_search_->Plan(start_int(0), start_int(1), goal_int(0),
                                   goal_int(1), search_region_);

  map_raw_path_ = graph_search_->GetPath();
  if (map_raw_path_.empty()) {
    if (planner_verbose_)
      std::cout << ANSI_COLOR_RED "Cannot find a path from " << start.transpose() << " to " << goal.transpose()
                << " Abort!" ANSI_COLOR_RESET << std::endl;
    status_ = -1;
    return false;
  }
//  //**** raw path, s --> g
  VectorPoint2f ps;
  for (const auto &point : map_raw_path_) {
    ps.push_back(map_util_->MapToWorld(point));
  }

  raw_path_.swap(ps);
  return true;
}

bool DMPlanner::Plan(const Point2f &start, const Point2f &goal, const VectorPoint2i &path) {
  if (planner_verbose_) {
    printf("****************[DistanceMapPlanner]***************\n");
    printf("eps: %f\n", eps_);
    printf("cweight: %f\n", cweight_);
    printf("pow: %d\n", pow_);
    std::cout << "search_radius: " << search_radius_.transpose() << std::endl;
    std::cout << "potential_radius: " << potential_radius_.transpose() << std::endl;
    std::cout << "potential_map_range: " << potential_map_range_.transpose() << std::endl;
  }

  search_region_ = SetPath(path, search_radius_, false); // sparse input path

  return Plan(start, goal, eps_, cweight_);
}

__attribute__((unused)) IterativeDMPlanner::IterativeDMPlanner(bool verbose) : DMPlanner(verbose) {
  if (this->planner_verbose_)
    printf(ANSI_COLOR_CYAN "Iterative DMP PLANNER VERBOSE ON\n" ANSI_COLOR_RESET);
}

__attribute__((unused)) bool IterativeDMPlanner::iterativeComputePath(
    const Point2f &start, const Point2f &goal,
    const VectorPoint2i &prior_path, int max_iteration) {
  if (this->planner_verbose_) {
    printf("****************[IterativeDistanceMapPlanner]***************\n");
    printf("eps: %f\n", this->eps_);
    printf("cweight: %f\n", this->cweight_);
    printf("pow: %d\n", this->pow_);
    printf("max_iteration: %d\n", max_iteration);
    std::cout << "search_radius: " << this->search_radius_.transpose() << std::endl;
    std::cout << "potential_radius: " << this->potential_radius_.transpose() << std::endl;
    std::cout << "potential_map_range: " << this->potential_map_range_.transpose() << std::endl;
  }

  VectorPoint2i path = prior_path;
  float prev_cost = std::numeric_limits<float>::infinity();
  for (int i = 0; i < max_iteration; i++) {
    this->search_region_ = this->SetPath(path, this->search_radius_, false); // sparse input path
    if (!this->Plan(start, goal, this->eps_, this->cweight_))
      return false;
    if (this->path_cost_ == prev_cost) {
      if (this->planner_verbose_)
        printf("[IterativeDMPlanner]: converged to %f!\n", prev_cost);
      break;
    } else {
      prev_cost = this->path_cost_;
      path = this->map_raw_path_;
      if (this->planner_verbose_)
        printf("[IterativeDMPlanner]: path cost is %f! iteration: %d\n",
               prev_cost, i);
    }
  }

  return true;
}

}


